#!/bin/bash

# PostgreSQL连接探测脚本
# 使用psql建立长连接，每2秒打印一次连接信息，输出到指定目录

# 默认配置
DEFAULT_HOST="localhost"
DEFAULT_PORT="5432"
DEFAULT_INTERVAL=2
DEFAULT_OUTPUT_DIR="./logs"

# 帮助信息
show_help() {
    cat << EOF
PostgreSQL连接探测脚本

用法: $0 -h HOST -p PORT -d DATABASE -u USER [-P PASSWORD] [-o OUTPUT_DIR] [-i INTERVAL]

参数:
    -h HOST         PostgreSQL主机地址 (必需)
    -p PORT         端口号 (默认: 5432)
    -d DATABASE     数据库名 (必需)
    -u USER         用户名 (必需)
    -P PASSWORD     密码 (可选，不提供时会提示输入)
    -o OUTPUT_DIR   日志输出目录 (默认: ./logs)
    -i INTERVAL     检测间隔秒数 (默认: 2)
    --help          显示此帮助信息

示例:
    $0 -h localhost -p 5432 -d mydb -u postgres -P mypass
    $0 -h 192.168.1.100 -d testdb -u admin -o /var/log/pg_probe
EOF
}

# 解析命令行参数
while [[ $# -gt 0 ]]; do
    case $1 in
        -h)
            PG_HOST="$2"
            shift 2
            ;;
        -p)
            PG_PORT="$2"
            shift 2
            ;;
        -d)
            PG_DATABASE="$2"
            shift 2
            ;;
        -u)
            PG_USER="$2"
            shift 2
            ;;
        -P)
            PGPASSWORD="$2"
            shift 2
            ;;
        -o)
            OUTPUT_DIR="$2"
            shift 2
            ;;
        -i)
            INTERVAL="$2"
            shift 2
            ;;
        --help)
            show_help
            exit 0
            ;;
        *)
            echo "未知参数: $1"
            show_help
            exit 1
            ;;
    esac
done

# 检查必需参数
if [[ -z "$PG_HOST" || -z "$PG_DATABASE" || -z "$PG_USER" ]]; then
    echo "错误: 缺少必需参数"
    show_help
    exit 1
fi

# 设置默认值
PG_PORT=${PG_PORT:-$DEFAULT_PORT}
INTERVAL=${INTERVAL:-$DEFAULT_INTERVAL}
OUTPUT_DIR=${OUTPUT_DIR:-$DEFAULT_OUTPUT_DIR}

# 如果没有提供密码，提示输入
if [[ -z "$PGPASSWORD" ]]; then
    echo -n "请输入密码: "
    read -s PGPASSWORD
    echo
fi

export PGPASSWORD

# 创建输出目录
mkdir -p "$OUTPUT_DIR"

# 生成日志文件名
TIMESTAMP=$(date '+%Y%m%d_%H%M%S')
LOG_FILE="$OUTPUT_DIR/pg_probe_${PG_HOST}_${PG_PORT}_${TIMESTAMP}.log"

# 日志函数
log_message() {
    local message="$1"
    local timestamp=$(date '+%Y-%m-%d %H:%M:%S.%3N')
    local formatted_message="[$timestamp] $message"
    echo "$formatted_message"
    echo "$formatted_message" >> "$LOG_FILE"
}

# 获取格式化时间戳
get_timestamp() {
    date '+%Y-%m-%d %H:%M:%S.%3N'
}

# 测试连接函数
test_connection() {
    psql -h "$PG_HOST" -p "$PG_PORT" -d "$PG_DATABASE" -U "$PG_USER" \
         -c "SELECT 1;" > /dev/null 2>&1
    return $?
}

# 获取连接信息
get_connection_info() {
    local query="
    SELECT 
        'PostgreSQL ' || version() as version,
        NOW() as server_time,
        (SELECT count(*) FROM pg_stat_activity) as total_connections,
        (SELECT count(*) FROM pg_stat_activity WHERE state = 'active') as active_connections,
        (SELECT count(*) FROM pg_stat_activity WHERE state = 'idle') as idle_connections,
        pg_size_pretty(pg_database_size('$PG_DATABASE')) as database_size;
    "
    
    psql -h "$PG_HOST" -p "$PG_PORT" -d "$PG_DATABASE" -U "$PG_USER" \
         -t -A -F'|' -c "$query" 2>/dev/null
}

# 打印连接信息
print_connection_info() {
    if test_connection; then
        local info=$(get_connection_info)
        if [[ -n "$info" ]]; then
            IFS='|' read -r version server_time total_conn active_conn idle_conn db_size <<< "$info"
            
            log_message "连接状态: ✓ CONNECTED | 服务器: ${PG_HOST}:${PG_PORT} | 数据库: ${PG_DATABASE}"
            log_message "版本信息: $version"
            log_message "连接统计: 总连接数=${total_conn} | 活跃=${active_conn} | 空闲=${idle_conn} | 数据库大小=${db_size}"
            log_message "服务器时间: $server_time"
        else
            log_message "连接状态: ✓ CONNECTED | 服务器: ${PG_HOST}:${PG_PORT} | 数据库: ${PG_DATABASE} | 获取详细信息失败"
        fi
    else
        log_message "连接状态: ✗ DISCONNECTED | 服务器: ${PG_HOST}:${PG_PORT} | 数据库: ${PG_DATABASE}"
    fi
}

# 信号处理函数
cleanup() {
    log_message "收到中断信号，正在关闭..."
    log_message "程序已退出"
    exit 0
}

# 设置信号处理
trap cleanup SIGINT SIGTERM

# 主程序
main() {
    log_message "PostgreSQL连接探测器启动"
    log_message "目标: ${PG_HOST}:${PG_PORT}/${PG_DATABASE}"
    log_message "用户: ${PG_USER}"
    log_message "检测间隔: ${INTERVAL}秒"
    log_message "日志文件: $LOG_FILE"
    log_message "$(printf '%80s' | tr ' ' '-')"
    
    # 初始连接测试
    if ! test_connection; then
        log_message "初始连接失败，请检查连接参数"
        exit 1
    fi
    
    log_message "初始连接成功，开始监控..."
    
    # 主循环
    while true; do
        print_connection_info
        log_message "$(printf '%50s' | tr ' ' '-')"
        sleep "$INTERVAL"
    done
}

# 启动主程序
main